# -*- coding: utf-8 -*-
"""
RandomForest股票预测模型预测脚本
"""

import sys
import os
import pandas as pd
import numpy as np
import joblib
import argparse
from datetime import datetime
import efinance as ef

def get_company_name(stock_code):
    """获取公司名称"""
    try:
        # 获取股票基本信息
        stock_info = ef.stock.get_base_info(stock_code)
        if stock_info and len(stock_info) > 0:
            return stock_info['股票名称'].iloc[0]
        return f"股票{stock_code}"
    except:
        return f"股票{stock_code}"

def predict_next_day(stock_code, frequency):
    """使用RandomForest模型预测下一个交易日"""
    
    print(f"🌳 RandomForest股票预测")
    print("=" * 50)
    
    # 文件名
    model_filename = f'output/models/rf_model_{stock_code}_{frequency}.pkl'
    data_filename = f'output/data/stock_data_{stock_code}_{frequency}_with_indicators.csv'
    
    try:
        # 获取公司名称
        company_name = get_company_name(stock_code)
        print(f"📊 股票信息: {stock_code} - {company_name}")
        print(f"⏰ 预测周期: {frequency}")
        print(f"📅 预测时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        
        # 加载模型
        print(f"\\n🤖 加载RandomForest模型: {model_filename}")
        if not os.path.exists(model_filename):
            print(f"❌ 模型文件不存在: {model_filename}")
            print(f"请先运行训练命令: python3 04c_train_randomforest.py {stock_code} {frequency}")
            return
        
        model = joblib.load(model_filename)
        print("✅ 模型加载成功")
        
        # 获取最新数据
        print("\\n📊 获取最新市场数据...")
        try:
            # 尝试从本地文件获取最新数据
            if os.path.exists(data_filename):
                df = pd.read_csv(data_filename)
                print("   使用本地数据文件")
            else:
                print("   本地数据不存在，需要先运行数据获取流程")
                return
            
            # 修复列名问题
            if 'date' in df.columns:
                df.rename(columns={'date': '日期'}, inplace=True)
            
            # 设置日期索引
            df['日期'] = pd.to_datetime(df['日期'])
            df.set_index('日期', inplace=True)
            
        except Exception as e:
            print(f"❌ 数据获取失败: {e}")
            return
        
        # 准备特征数据
        feature_columns = [col for col in df.columns if col not in ['code', 'close', 'target']]
        latest_features = df[feature_columns].iloc[-1:].values
        
        print(f"   特征维度: {latest_features.shape}")
        print(f"   数据日期: {df.index[-1].strftime('%Y-%m-%d')}")
        
        # 进行预测
        print("\\n🔮 模型预测中...")
        predicted_price = model.predict(latest_features)[0]
        
        # 获取当前价格信息
        current_data = df.iloc[-1]
        current_price = current_data['close']
        current_open = current_data['open']
        current_high = current_data['high']
        current_low = current_data['low']
        current_volume = current_data['volume']
        
        # 计算预测变化
        price_change = predicted_price - current_price
        percent_change = (price_change / current_price) * 100
        
        # 显示预测结果
        print("\\n" + "="*60)
        print("🔮 RandomForest 预测结果")
        print("="*60)
        print(f"📊 股票代码: {stock_code} ({company_name})")
        print(f"⏰ 预测周期: {frequency}")
        print(f"📅 数据基准: {df.index[-1].strftime('%Y-%m-%d')}")
        print("-"*60)
        print(f"💰 当前收盘价: {current_price:.2f}元")
        print(f"📈 当日开盘价: {current_open:.2f}元")
        print(f"📈 当日最高价: {current_high:.2f}元")
        print(f"📉 当日最低价: {current_low:.2f}元")
        print(f"💾 当日成交量: {current_volume:,.0f}手")
        print("-"*60)
        print(f"🔮 预测收盘价: {predicted_price:.2f}元")
        print(f"📊 预测涨跌: {price_change:+.2f}元 ({percent_change:+.2f}%)")
        
        # 趋势判断
        if percent_change > 2:
            trend = "📈 强烈看涨"
            color = "🟢"
        elif percent_change > 0.5:
            trend = "📈 看涨"
            color = "🟢"
        elif percent_change > -0.5:
            trend = "➡️ 震荡"
            color = "🟡"
        elif percent_change > -2:
            trend = "📉 看跌"
            color = "🔴"
        else:
            trend = "📉 强烈看跌"
            color = "🔴"
        
        print(f"📈 预测趋势: {color} {trend}")
        print("="*60)
        
        # 模型置信度信息
        print("\\n📊 模型信息:")
        print(f"   模型类型: RandomForest")
        print(f"   树的数量: {model.n_estimators}")
        print(f"   最大深度: {model.max_depth}")
        if hasattr(model, 'oob_score_'):
            print(f"   袋外得分: {model.oob_score_:.4f}")
        
        # 风险提示
        print("\\n⚠️ 风险提示:")
        print("   • 本预测仅供参考，不构成投资建议")
        print("   • 股市有风险，投资需谨慎")
        print("   • 请结合其他分析方法综合判断")
        
        print("\\n🚀 RandomForest 预测流程完成！ 🚀")
        
    except Exception as e:
        print(f"❌ 预测失败: {e}")
        import traceback
        traceback.print_exc()

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='RandomForest股票预测')
    parser.add_argument('stock_code', help='股票代码')
    parser.add_argument('frequency', help='数据频率 (daily/hourly)')
    
    args = parser.parse_args()
    
    predict_next_day(args.stock_code, args.frequency)

if __name__ == "__main__":
    main()